{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# 超参数设置\n",
    "batch_size = 64\n",
    "epochs = 5\n",
    "lr = 0.001\n",
    "\n",
    "# 数据加载（MNIST）\n",
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]\n",
    ")\n",
    "train_dataset = datasets.MNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform\n",
    ")\n",
    "test_dataset = datasets.MNIST(root=\"./data\", train=False, transform=transform)\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义神经网络结构\n",
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(784, 128)\n",
    "        self.fc2 = nn.Linear(128, 64)\n",
    "        self.fc3 = nn.Linear(64, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(-1, 784)  # 展平图像\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = F.relu(self.fc2(x))\n",
    "        x = self.fc3(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== CrossEntropy Loss ===\n",
      "Epoch 1 [CE]: Loss=0.3927, Acc=88.08%\n",
      "Epoch 2 [CE]: Loss=0.1827, Acc=94.55%\n",
      "Epoch 3 [CE]: Loss=0.1335, Acc=95.88%\n",
      "Epoch 4 [CE]: Loss=0.1077, Acc=96.59%\n",
      "Epoch 5 [CE]: Loss=0.0922, Acc=97.10%\n",
      "\n",
      "=== MSE Loss ===\n",
      "Epoch 1 [MSE]: Loss=0.0178, Acc=88.16%\n",
      "Epoch 2 [MSE]: Loss=0.0095, Acc=93.79%\n",
      "Epoch 3 [MSE]: Loss=0.0072, Acc=95.27%\n",
      "Epoch 4 [MSE]: Loss=0.0059, Acc=96.18%\n",
      "Epoch 5 [MSE]: Loss=0.0054, Acc=96.55%\n"
     ]
    }
   ],
   "source": [
    "# 对比实验函数\n",
    "def train_with_loss(loss_type=\"CE\"):\n",
    "    model = Net()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "\n",
    "    if loss_type == \"CE\":\n",
    "        criterion = nn.CrossEntropyLoss()  # 内置Softmax\n",
    "    else:\n",
    "        criterion = nn.MSELoss()  # 需要手动Softmax\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        total_loss, correct = 0, 0\n",
    "\n",
    "        for images, labels in train_loader:\n",
    "            optimizer.zero_grad()\n",
    "            outputs = model(images)\n",
    "\n",
    "            # MSE需要特殊处理\n",
    "            if loss_type == \"MSE\":\n",
    "                # 将标签转为one-hot编码\n",
    "                targets = F.one_hot(labels, num_classes=10).float()\n",
    "                # 对输出应用Softmax\n",
    "                outputs = F.softmax(outputs, dim=1)\n",
    "                loss = criterion(outputs, targets)\n",
    "            else:\n",
    "                loss = criterion(outputs, labels)\n",
    "\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            total_loss += loss.item()\n",
    "            preds = (\n",
    "                outputs.argmax(dim=1) if loss_type == \"CE\" else outputs.argmax(dim=1)\n",
    "            )\n",
    "            correct += (preds == labels).sum().item()\n",
    "\n",
    "        # 打印训练结果\n",
    "        train_acc = 100 * correct / len(train_dataset)\n",
    "        print(\n",
    "            f\"Epoch {epoch+1} [{loss_type}]: Loss={total_loss/len(train_loader):.4f}, Acc={train_acc:.2f}%\"\n",
    "        )\n",
    "\n",
    "\n",
    "# 运行对比实验\n",
    "print(\"=== CrossEntropy Loss ===\")\n",
    "train_with_loss(\"CE\")\n",
    "\n",
    "print(\"\\n=== MSE Loss ===\")\n",
    "train_with_loss(\"MSE\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mytorchenv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
